#include "ForkPreamble.h"
#include <signal.h>
#include "Slicer.h"

map<Function*, StubGenerator*> StubGenerator::ByDst;

StubGenerator::StubGenerator(Function *F, Instruction *&InsertBefore, ExecutionEngine *EE)
: C(F->getContext()), F(F), M(F->getParent()), EE(EE)
{
	Slice = GenerateEmptySlice(F);
	Stub = Function::Create(F->getFunctionType(), Function::InternalLinkage, "", M);

	BasicBlock *EntryBB = BasicBlock::Create(C, "entry", Stub);
	BasicBlock *ForkBB = BasicBlock::Create(C, "fork", Stub);
	BasicBlock *CallBB = BasicBlock::Create(C, "call", Stub);
	BasicBlock *RetBB = BasicBlock::Create(C, "last", Stub);

	const Type *Int32Ty = Type::getInt32Ty(C);
	Constant *Int0 = ConstantInt::get(Int32Ty, 0);
	Pid = new GlobalVariable(*M, Int32Ty, false, GlobalValue::InternalLinkage, Int0, "");
	ShouldFork = new GlobalVariable(*M, Int32Ty, false, GlobalValue::InternalLinkage, Int0, "");

	IRBuilder<> B(C);

	B.SetInsertPoint(EntryBB);
	B.CreateCondBr(B.CreateICmpEQ(B.CreateLoad(ShouldFork), Int0), RetBB, ForkBB);

	B.SetInsertPoint(ForkBB);
	Value *ForkRetVal = B.CreateCall(M->getFunction("fork"));
	B.CreateCondBr(B.CreateICmpEQ(ForkRetVal, Int0), CallBB, RetBB);

	B.SetInsertPoint(CallBB);
	B.CreateStore(ForkRetVal, Pid);
	std::vector<Value*> Args;
	for (Function::arg_iterator AI = Stub->arg_begin(), AE = Stub->arg_end(); AI != AE; ++AI)
		Args.push_back(&*AI);
//	B.CreateCall(M->getFunction("_log"), ConstantInt::get(Type::getInt32Ty(C), 8));
	B.CreateCall(Slice, Args.begin(), Args.end());
//	B.CreateCall(M->getFunction("_log"), ConstantInt::get(Type::getInt32Ty(C), 9));
	B.CreateCall(M->getFunction("exit"), Int0);
	B.CreateUnreachable();

	B.SetInsertPoint(RetBB);
	if (Stub->getReturnType()->isVoidTy())
		B.CreateRetVoid();
	else
		B.CreateRet(Constant::getNullValue(Stub->getReturnType()));

	verifyFunction(*Stub, llvm::AbortProcessAction);

	B.SetInsertPoint(InsertBefore);

	vector<Value*> Args2;
	for (Function::arg_iterator AI = F->arg_begin(), AE = F->arg_end(); AI != AE; ++AI)
		Args2.push_back(&*AI);
	InsertBefore = B.CreateCall(Stub, Args2.begin(), Args2.end());

	PidPtr = (int32_t*)EE->getPointerToGlobal(Pid);
	ShouldForkPtr = (int32_t*)EE->getPointerToGlobal(ShouldFork);
	*PidPtr = 0;

	ByDst[Slice] = this;
}

Function* StubGenerator::GenerateEmptySlice(Function *F)
{
	Function *NF = Function::Create(F->getFunctionType(), Function::InternalLinkage, "", M);

	IRBuilder<> Builder(BasicBlock::Create(C, "entry", NF));
	if (NF->getReturnType()->isVoidTy())
		Builder.CreateRetVoid();
	else
		Builder.CreateRet(Constant::getNullValue(NF->getReturnType()));

	return NF;
}

bool StubGenerator::SliceAndCompile(vector<Instruction*> &HotLoads)
{
	*PidPtr = 0;
	*ShouldForkPtr = 0;
	ReverseVMap.clear();
	vector<Instruction*> Splitters;
	set<Instruction*> ActualLoads;

	for (vector<Instruction*>::iterator I = HotLoads.begin(), E = HotLoads.end(); I != E; ++I)
		ActualLoads.insert(*I);

	if (ActualLoads.empty())
				return false;

	errs() << "\nSlicing\n";

	for (;;) {
		if (ActualLoads.empty())
			return false;

		errs() << "\tRound\n";

		errs() << "\tHotLoads:\n";
		for (set<Instruction*>::iterator I = ActualLoads.begin(), E = ActualLoads.end(); I != E; ++I)
			errs() << "\t" << **I << "\n";

		errs() << "\tSplitter:\n";
		for (vector<Instruction*>::iterator I = Splitters.begin(), E = Splitters.end(); I != E; ++I)
			errs() << "\t" << **I << "\n";

		errs() << "\n";

		Slicer FuncSlicer(F, Slice, Splitters, EE, ReverseVMap);
		Instruction* Inst = FuncSlicer.Slice(ActualLoads);
		if (!Inst)
			break;
		ActualLoads.erase(Inst);
		Splitters.push_back(Inst);
	}
	*ShouldForkPtr = 1;
	EE->recompileAndRelinkFunction(Slice);
	return true;
}

void StubGenerator::Kill()
{
	if (*PidPtr) {
		errs() << "killing " << *PidPtr << "\n";
		kill(*PidPtr, 9);
	}
}

